from functools import wraps
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from starlette.requests import Request

SQLALCHEMY_DATABASE_URL = "mysql+pymysql://root:admin1234@sh-cdb-k4k80j2k.sql.tencentcdb.com:63742/studytest"
# SQLALCHEMY_DATABASE_URL = "postgresql://user:password@postgresserver/db"

engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    # True开启调试
    echo=True,
    pool_pre_ping=True,
    pool_size=5,
    max_overflow=10
)

SessionLocal = sessionmaker(
    autocommit=False,
    # class_=AsyncSession,
    autoflush=False,
    expire_on_commit=False,
    bind=engine
)

Base = declarative_base()


# 定义一个闭包，用于为 API 接口函数增加数据库异常回滚的功能
def db_rollback_on_error(func):
    @wraps(func)
    def wrapper(req: Request, *args, **kwargs):
        db = SessionLocal()
        req.state.db = db
        try:
            response = func(req, *args, **kwargs)
        except Exception as e:

            db.rollback()
            raise
        else:
            db.commit()
        finally:
            db.close()
        return response

    return wrapper
